import os
from typing import Optional, List

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

os.environ['LANGCHAIN_TRACING_V2'] = "true"
os.environ['LANGCHAIN_API_KEY'] = '1123'

# 调用大模型

model = ChatOpenAI(model='gpt-4-turbo')


# pydantic:处理数据，验证数据，定义数据的格式，虚拟化和反虚拟化，类型装欢等等
class Person(BaseModel):
    # 关于一个人的数据模型
    name: Optional[str] = Field(default=None, description='表示人的名字')
    hair_color: Optional[str] = Field(default=None, description='如果知道的话，这个人的头发颜色')
    height_in_meters: Optional[str] = Field(default=None, description='以米为单位测量的高度')


# 不用manyPerson，只会一直返回一条数据
class ManyPerson(BaseModel):
    # 关于一个人的数据模型
    people: List[Person]


prompt = ChatPromptTemplate.from_messages([
    {
        "system", "你是一个专业的抽取算法。只从未结构化文本中提取相关信息。如果你不知道要提取的属性的值，返回该属性值为null"
    },
    # MessagePlaceholder('examples')
    ("human", "{text}")
])

chain = prompt | {'text': RunnablePassthrough()} | model.with_structured_output(schema=ManyPerson)

resp = chain.invoke('马路上走来一个女生，长长的头发披在肩上，大概1米7')

print(resp)
